import math
import struct

# 全局变量
q = [1.0, 0.0, 0.0, 0.0]  # 四元数 [w, x, y, z]
P = [[0.1, 0, 0, 0],      # 协方差矩阵
     [0, 0.1, 0, 0],
     [0, 0, 0.1, 0],
     [0, 0, 0, 0.1]]

# 滤波参数
Kp = 0.8     # 互补滤波比例系数
Ki = 0.001   # 互补滤波积分系数
Q = 0.001    # 卡尔曼滤波过程噪声
R = 0.1      # 卡尔曼滤波测量噪声
dt = 0.005   # 采样时间

# 角速度积分
wx_int = 0.0
wy_int = 0.0
wz_int = 0.0

def Q_rsqrt(number):
    """
    快速计算1/sqrt(number)
    """
    threehalfs = 1.5
    x2 = number * 0.5
    y = number
    
    # 将浮点数转换为整数
    i = struct.unpack('I', struct.pack('f', y))[0]
    i = 0x5f3759df - (i >> 1)
    y = struct.unpack('f', struct.pack('I', i))[0]
    
    # 牛顿迭代
    y = y * (threehalfs - (x2 * y * y))
    return y

def quaternion_multiply(q1, q2):
    """
    四元数乘法
    """
    w1, x1, y1, z1 = q1
    w2, x2, y2, z2 = q2
    
    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
    
    return [w, x, y, z]

def quaternion_normalize(q):
    """
    四元数归一化
    """
    norm = math.sqrt(q[0]*q[0] + q[1]*q[1] + q[2]*q[2] + q[3]*q[3])
    if norm > 0:
        return [q[0]/norm, q[1]/norm, q[2]/norm, q[3]/norm]
    return q

def process_imu_data(imu_data):
    """
    处理原始IMU数据
    """
    # 加速度计数据转换 (单位: m/s^2)
    ax = imu_data[0] * 0.244 / 1000
    ay = imu_data[1] * 0.244 / 1000
    az = imu_data[2] * 0.244 / 1000
    
    # 陀螺仪数据转换 (单位: rad/s)
    gx = imu_data[3] * 0.07 * 0.0174
    gy = imu_data[4] * 0.07 * 0.0174
    gz = imu_data[5] * 0.07 * 0.0174
    
    return (ax, ay, az, gx, gy, gz)

def complementary_filter(gyro_data, accel_data):
    """
    互补滤波更新四元数
    """
    global q, wx_int, wy_int, wz_int
    gx, gy, gz = gyro_data
    ax, ay, az = accel_data
    
    # 归一化加速度数据
    norm = math.sqrt(ax*ax + ay*ay + az*az)
    if norm > 0:
        ax /= norm
        ay /= norm
        az /= norm
    
    # 计算加速度计姿态
    accel_pitch = math.asin(max(min(ax, 1.0), -1.0))
    accel_roll = math.atan2(ay, az)
    
    # 从四元数计算当前姿态
    sin_pitch = 2 * (q[0]*q[2] - q[1]*q[3])
    sin_pitch = max(min(sin_pitch, 1.0), -1.0)
    current_pitch = math.asin(sin_pitch)
    
    sin_roll = 2 * (q[0]*q[1] + q[2]*q[3])
    cos_roll = 1 - 2 * (q[1]*q[1] + q[2]*q[2])
    current_roll = math.atan2(sin_roll, cos_roll)
    
    # 计算姿态误差
    pitch_error = accel_pitch - current_pitch
    roll_error = accel_roll - current_roll
    
    # 更新角速度积分
    wx_int += pitch_error * Ki
    wy_int += roll_error * Ki
    
    # 限制积分范围
    wx_int = max(min(wx_int, 0.5), -0.5)
    wy_int = max(min(wy_int, 0.5), -0.5)
    
    # 更新角速度
    gx += Kp * pitch_error + wx_int
    gy += Kp * roll_error + wy_int
    
    # 更新四元数
    q_dot = [0.5 * (-q[1]*gx - q[2]*gy - q[3]*gz),
             0.5 * (q[0]*gx + q[2]*gz - q[3]*gy),
             0.5 * (q[0]*gy - q[1]*gz + q[3]*gx),
             0.5 * (q[0]*gz + q[1]*gy - q[2]*gx)]
    
    q = [q[i] + q_dot[i] * dt for i in range(4)]
    q = quaternion_normalize(q)

def kalman_filter(gyro_data, accel_data):
    """
    卡尔曼滤波更新四元数
    """
    global q, P
    gx, gy, gz = gyro_data
    ax, ay, az = accel_data
    
    # 预测步骤
    q_dot = [0.5 * (-q[1]*gx - q[2]*gy - q[3]*gz),
             0.5 * (q[0]*gx + q[2]*gz - q[3]*gy),
             0.5 * (q[0]*gy - q[1]*gz + q[3]*gx),
             0.5 * (q[0]*gz + q[1]*gy - q[2]*gx)]
    
    q = [q[i] + q_dot[i] * dt for i in range(4)]
    q = quaternion_normalize(q)
    
    # 更新协方差矩阵
    for i in range(4):
        for j in range(4):
            P[i][j] += Q * dt
    
    # 归一化加速度数据
    norm = math.sqrt(ax*ax + ay*ay + az*az)
    if norm > 0:
        ax /= norm
        ay /= norm
        az /= norm
    
    # 构建测量矩阵
    H = [[0, 2*q[2], -2*q[3], 2*q[0]],
         [-2*q[1], 0, 2*q[0], 2*q[3]],
         [2*q[0], 2*q[1], 2*q[2], 0]]
    
    # 计算测量残差
    h = [2*(q[1]*q[3] - q[0]*q[2]),
         2*(q[0]*q[1] + q[2]*q[3]),
         q[0]*q[0] - q[1]*q[1] - q[2]*q[2] + q[3]*q[3]]
    
    y = [ax - h[0], ay - h[1], az - h[2]]
    
    # 计算卡尔曼增益
    S = [[0.0 for _ in range(3)] for _ in range(3)]
    for i in range(3):
        for j in range(3):
            S[i][j] = R if i == j else 0.0
            for k in range(4):
                S[i][j] += H[i][k] * P[k][j]
    
    K = [[0.0 for _ in range(3)] for _ in range(4)]
    for i in range(4):
        for j in range(3):
            K[i][j] = 0.0
            for k in range(4):
                K[i][j] += P[i][k] * H[j][k]
            K[i][j] /= S[j][j]
    
    # 更新状态
    for i in range(4):
        for j in range(3):
            q[i] += K[i][j] * y[j]
    
    q = quaternion_normalize(q)
    
    # 更新协方差矩阵
    for i in range(4):
        for j in range(4):
            P[i][j] -= sum(K[i][k] * S[k][j] for k in range(3))

def get_angles():
    """
    从四元数计算欧拉角
    """
    # 计算俯仰角
    sin_pitch = 2 * (q[0]*q[2] - q[1]*q[3])
    sin_pitch = max(min(sin_pitch, 1.0), -1.0)
    pitch = math.asin(sin_pitch) * 57.3
    
    # 计算横滚角
    sin_roll = 2 * (q[0]*q[1] + q[2]*q[3])
    cos_roll = 1 - 2 * (q[1]*q[1] + q[2]*q[2])
    roll = math.atan2(sin_roll, cos_roll) * 57.3
    
    # 计算偏航角
    sin_yaw = 2 * (q[0]*q[3] + q[1]*q[2])
    cos_yaw = 1 - 2 * (q[2]*q[2] + q[3]*q[3])
    yaw = math.atan2(sin_yaw, cos_yaw) * 57.3
    
    return pitch, roll, yaw

def get_accel_angles(accel_data):
    """
    使用加速度计数据直接计算姿态角
    """
    ax, ay, az = accel_data
    
    # 归一化加速度数据
    norm = math.sqrt(ax*ax + ay*ay + az*az)
    if norm > 0:
        ax /= norm
        ay /= norm
        az /= norm
    
    # 确保输入值在[-1, 1]范围内
    ax = max(min(ax, 1.0), -1.0)
    
    # 计算俯仰角和横滚角
    pitch = math.asin(ax) * 57.3
    roll = math.atan2(ay, az) * 57.3
    
    return pitch, roll

def fusion_angle_cal(imu_data):
    """
    融合姿态解算主函数
    """
    # 处理IMU数据
    ax, ay, az, gx, gy, gz = process_imu_data(imu_data)
    
    # 互补滤波更新
    complementary_filter((gx, gy, gz), (ax, ay, az))
    
    # 卡尔曼滤波更新
    kalman_filter((gx, gy, gz), (ax, ay, az))
    
    # 获取姿态角
    pitch, roll, yaw = get_angles()
    
    # 获取加速度计直接解算的姿态角
    accel_pitch, accel_roll = get_accel_angles((ax, ay, az))
    
    return (pitch, roll, yaw), (accel_pitch, accel_roll) 